from scipy.ndimage.morphology import binary_dilation
import librosa
import numpy as np
import struct
import librosa.display
import webrtcvad
import soundfile as sf
import os

int16_max = (2 ** 15) - 1
#输入
# filename ='/Users/qianghaozhang/Desktop/科沃斯/智能语音数据/测试集/YIKO海外测试集/UK(1-14)/UK-male-1.wav'
filename="/Users/qianghaozhang/Desktop/科沃斯/智能语音数据/录音文件/1/1.1.1.wav"
# save_file = "/Users/qianghaozhang/Desktop/科沃斯/智能语音数据/测试集/YIKO海外测试集/UK(1-14)/"
save_file="/Users/qianghaozhang/Desktop/科沃斯/智能语音数据/录音文件/1/"
wav, source_sr = librosa.load(filename, sr=16000)

#计算语音检测窗口大小  //为整除 30秒X16000=总帧长#
samples_per_window = (10 * 16000) // 1000

#修剪音频的结尾，使其具有窗口大小的倍数。使wav的长度能被 samples_per_window整除
wav = wav[:len(wav) - (len(wav) % samples_per_window)]

pcm_wave = struct.pack("%dh" % len(wav), *(np.round(wav * int16_max)).astype(np.int16))

# 浮点数波形转换为16位单声道PCM  *：接收到的参数会形成一个元组，**：接收到的参数会形成一个字典。如下代码。
# webrtcvad 的 is_speech 接收的是buf 所以这里需要转换
pcm_wave = struct.pack("%dh" % len(wav), *(np.round(wav * int16_max)).astype(np.int16))

# 这里共有三种帧长可以用到，分别是80/10ms，160/20ms，240/30ms。其它采样率
# 的48k，32k，24k，16k会重采样到8k来计算VAD。之所以选择上述三种帧长度，是因为语
# // 音信号是短时平稳信号，其在10ms~30ms之间可看成平稳信号，高斯马尔科夫等比较
# // 的信号处理方法基于的前提是信号是平稳的，在10ms~30ms，平稳信号处理方法是可
# // 以使用的。
# // 　　从vad的代码中可以看出，实际上，系统只处理默认10ms,20ms,30ms长度的数据，
# // 其它长度的数据没有支持，笔者修改过可以支持其它在10ms-30ms之间长度的帧长度
# // 发现也是可以的。
# // 　　vad检测共四种模式，用数字0~3来区分，激进程度与数值大小正相关。
# // 0: Normal，1：low Bitrate， 2：Aggressive；3：Very Aggressive 可以根据实际的使用

voice_flags = []
vad = webrtcvad.Vad(mode=3)
for window_start in range(0, len(wav), samples_per_window):
    window_end = window_start + samples_per_window
    # append 进来的都是Boolean  这里以samples_per_windowx2 的长度去检测是否为人声
    voice_flags.append(vad.is_speech(pcm_wave[window_start * 2:window_end * 2],
                                     sample_rate=16000))



voice_flags = np.array(voice_flags)

def moving_average(array, width):
    #  拼接 bool 二值化
    # width 执行滑动平均平滑时，帧的平均数。
    # 该值越大，VAD变化必须越大才能平滑。
    array_padded = np.concatenate((np.zeros((width - 1) // 2), array, np.zeros(width // 2)))
    # 一维数组累加
    ret = np.cumsum(array_padded, dtype=float)
    ret[width:] = ret[width:] - ret[:-width]
    return ret[width - 1:] / width

audio_mask = moving_average(voice_flags, 8)
audio_mask = np.round(audio_mask).astype(np.bool)

# 扩张浊音区 使用多维二元膨胀 是数学形态学的方法 类似opencv 也有开闭运算 腐蚀膨胀
audio_mask = binary_dilation(audio_mask, np.ones(6 + 1))
audio_mask = np.repeat(audio_mask, samples_per_window)
res=wav[audio_mask == True]

def seperately_save():
    all_voice_lists = []
    tmp = []
    for flag,voice in zip(audio_mask,wav):
        if flag:
            tmp.append(voice)

        else:
            all_voice_lists.append(tmp)
            tmp=[]
    remove_silence=[]
    for voice_list in all_voice_lists:
        if voice_list !=[]:
            remove_silence.append(voice_list)

    for idx,speech in enumerate(remove_silence):
        speech = np.array(speech)
        path = os.path.join(save_file,"{}.wav".format(idx))
        sf.write(path,speech.astype(np.float32),16000,subtype="PCM_24")

if __name__ == '__main__':
    seperately_save()

